from langchain.chat_models import init_chat_model
import os
from typing import Optional
from pydantic import BaseModel, Field

key = os.getenv("OPENAI_API_KEY")
# print(key)
api_key = str(key)

llm = init_chat_model(
    model="gpt-4o-mini",
    model_provider="openai",
    base_url="https://api.zetatechs.com/v1",
    api_key=api_key
)


# 1. 方式为: Pydantic
# class Joke(BaseModel):
#     """Joke to tell user."""
#
#     setup: str = Field(description="The setup of the joke")
#     punchline: str = Field(description="The punchline to the joke")
#     rating: Optional[int] = Field(
#         default=None, description="How funny the joke is, from 1 to 10"
#     )
#
#
# structured_llm = llm.with_structured_output(Joke)
# print(structured_llm.invoke("Tell me a joke about cats"))
from typing import Optional
from typing_extensions import Annotated, TypedDict


# TypedDict
class Joke(TypedDict):
    """Joke to tell user."""

    setup: Annotated[str, ..., "The setup of the joke"]

    # Alternatively, we could have specified setup as:
    # setup: str                    # no default, no description
    # setup: Annotated[str, ...]    # no default, no description
    # setup: Annotated[str, "foo"]  # default, no description

    punchline: Annotated[str, ..., "The punchline of the joke"]
    rating: Annotated[Optional[int], None, "How funny the joke is, from 1 to 10"]

structured_llm = llm.with_structured_output(Joke)
print(structured_llm.invoke("Tell me a joke about cats"))